/* 
   fastaIO - IO for fasta records

   Copyright (C) 2004 Liu Tao

   This library is free software; you can redistribute it and/or
   modify it under the terms of the GNU Library General Public License
   as published by the Free Software Foundation; either version 2 of
   the License, or (at your option) any later version.

   This library is distributed in the hope that it will be useful, but
   WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
   Library General Public License for more details.

   You should have received a copy of the GNU Library General Public
   License along with this library; see the file COPYING.LIB.  If not,
   write to the Free Software Foundation, Inc., 59 Temple Place -
   Suite 330, Boston, MA 02111-1307, USA.

   #include <termios.h>
   #include <grp.h>
   #include <pwd.h>
*/

/* 
   RCS 
   $Id: fastaIO.c,v 1.3 2006/01/30 16:27:35 fred Exp $
   
   $Log: fastaIO.c,v $
   Revision 1.3  2006/01/30 16:27:35  fred
   Read and encode the incoming sequences. Guess the sequence type from
   the first 100 bp.

   Revision 1.2  2006/01/13 14:29:28  fred
   add revcomp and splice_c_record functions.

   Revision 1.1  2006/01/12 14:49:15  fred
   initiate import.


*/

#include <stdio.h>
#include <sys/types.h>
#include "system.h"
#include "taostr.h"
#include "fastaIO.h"

static void remove_enter ( const char *seq );

const char * SYMBOLTABLE = " ABCDEFGHIJKLMNOPQRSTUVWXYZ*";
const char * aminos      = "ABCDEFGHIKLMNPQRSTVWXYZ*";
const char * primenuc    = "ACGTU";


/* parse a file */
FastaRecord*
parse_file (FILE * fp)
{
  FastaRecord * records;
  int char_num;
  char *text;
  fseek(fp,0,SEEK_END);
  char_num = ftell(fp);
  rewind(fp);
  text = (char*) malloc(sizeof(char)*(char_num+1));
  fread(text,sizeof(char),char_num,fp);
  /* parse the whole text, produce a chain of fasta records */
  records = parse_text (text);
  free(text);
  return records;
}


/* parse whole fasta file */
FastaRecord* 
parse_text ( const char* text)
{
  char* ptr;
  char* record_text;
  FastaRecord *first_record,*cur_record,*next_record;

  ptr = (char*)text;

  record_text = strtok(ptr,">"); /* first record */
  first_record = parse_record(record_text);
  cur_record = first_record;

  while ( (record_text=strtok(NULL,">")) != NULL )
    {

      next_record = parse_record(record_text);
      cur_record->next = next_record;
      cur_record = next_record;
    }
  return first_record;
}

/* parse record text. */
FastaRecord*
parse_record ( const char* text)
{
  char *record_text;
  char *id,*seq,*ptr_first_enter;
  long length;
  FastaRecord *trecord;
  Fasta *tfasta;

  record_text = (char*)text;
  /* extract seq part. */
  ptr_first_enter = strchr(record_text,'\n');
  seq = ++ptr_first_enter;
  remove_enter(seq);

  /* extract id part */
  *(--ptr_first_enter) = '\0';
  id = record_text;

  length = strlen(seq);
  tfasta = (Fasta*)malloc(sizeof(Fasta));
  tfasta->id = (char*)strdup(id);
  tfasta->seq = (char*)strdup(seq);
  tfasta->length = length;
  
  trecord = (FastaRecord*) malloc(sizeof(FastaRecord));
  trecord->fasta = tfasta;
  trecord->fasta->stype = guess_type(trecord);
  trecord->next = NULL;

  return trecord;
}

/* create a fasta record */
FastaRecord* 
new_record  (const char* id, const char* seq )
{
  FastaRecord *record = (FastaRecord *)malloc(sizeof(FastaRecord));
  Fasta *fasta = (Fasta *)malloc(sizeof(Fasta));
  fasta->seq = (char*)strdup((char *)seq);
  fasta->length = strlen(seq);
  fasta->id = (char *)strdup(id);
  record->fasta = fasta;
  record->fasta->stype = guess_type(record);
  record->next = NULL;
  return record;
}

/* duplicate a fasta record */
FastaRecord* 
dup_record (const FastaRecord * src_record)
{
  FastaRecord *trgt_record = (FastaRecord *)malloc(sizeof(FastaRecord));
  Fasta *trgt_fasta = (Fasta *)malloc(sizeof(Fasta));
  trgt_fasta->seq = (char *)strdup(src_record->fasta->seq);
  trgt_fasta->length = src_record->fasta->length;
  trgt_fasta->id = (char *)strdup(src_record->fasta->id);
  trgt_record->fasta = trgt_fasta;
  trgt_record->next = NULL;
  return trgt_record;
}

/* guess type of a sequence from the first 100 bp */
char
guess_type ( const FastaRecord * record )
{
  char part[101];
  char * p;
  int num_A = 0, num_C = 0;
  float fre_A, fre_C;
  strncpy (part, record->fasta->seq, 100);
  part[100] = 0;
  /* if occurence frequency of A > 10% or C > 10%, then seq is a
	 nucleotide sequence.*/

  for (p=part;*p;p++)
	if ( *p == 'A' || *p == 'a' )
	  num_A ++;
	else if ( *p == 'C' || *p == 'c' )
	  num_C ++;
  fre_A = ((float)num_A) / strlen(part);
  fre_C = ((float)num_C) / strlen(part);

  if ( fre_A > 0.1 && fre_C > 0.1 )
	if ( index(part, 'U') )
	  return 'R';
	else
	  return 'D';
  else 
	return 'P';
}

/* uppercase letters in fasta sequence */
int 
up_record (const FastaRecord * record)
{
  char *p = record->fasta->seq;
  for ( ; *p ; p++ )
	*p = toupper(*p);
  return 1;
}

/* lowercase letters in fasta sequence */
int 
low_record (const FastaRecord * record)
{
  char *p = record->fasta->seq;
  for ( ; *p ; p++ )
	*p = tolower(*p);
  return 1;
}

/* get a subrecord [start,end), and `start' is from 0*/
FastaRecord* 
sub_record (const FastaRecord * record, int start, int end)
{
  FastaRecord *subrecord = (FastaRecord *)malloc(sizeof(FastaRecord));
  Fasta *subfasta = (Fasta *)malloc(sizeof(Fasta));
  char *seq = (char *)malloc(sizeof(char)*(end-start+1));
  int i;
  for (i=start; i<end; i++)
	seq[i-start] = record->fasta->seq[i];
  seq[end-start] = '\0';
  subfasta->seq = seq;
  subfasta->length = end-start;
  subfasta->id = (char *)strdup(record->fasta->id);
  subfasta->stype = record->fasta->stype;
  subrecord->fasta = subfasta;
  subrecord->next = NULL;
  return subrecord;
}

/* mask a part of record 
*/
int 
mask_record (const FastaRecord *record, int start, int end)
{
  int i;
  char mask;
  char *seq = record->fasta->seq;
  if (record->fasta->stype == 'R' || record->fasta->stype == 'D' )
	mask = 'N';
  else if (record->fasta->stype == 'P')
	mask = 'X';
  else
	return 0;
  for ( i=start; i<end; i++ )
	seq[i] = mask;
  return 1;
}

/* replace char X with char Y
*/
int 
replace_record (const FastaRecord *record, char X, char Y)
{
  char *p = record->fasta->seq;
  char x;
  char y;
  x = X;
  y = Y;
  while ( (p = index(p,x)) )
	*p = y;
  return 1;
}


/* write a record */
int 
write_record (FILE *fp, FastaRecord* fr)
{
  int lines,i,j,s,e;
  fprintf (fp, ">%s\n", fr->fasta->id);
  lines = (fr->fasta->length-1)/50+1;
  for ( i = 0; i< lines; i++ )
	{
	  s = i*50;
	  e = (i+1)*50;
	  if (e > fr->fasta->length) e = fr->fasta->length;
	  for ( j = s; j< e; j++)
		fputc( fr->fasta->seq[j], fp);
	  fputc ('\n',fp);
	}
  fputc ('\n',fp);
  return 1;
}

/* write records */
int
write_records (FILE *fp, FastaRecord* fr)
{
  FastaRecord *ptr;
  ptr = fr;
  while (ptr != NULL) 
    {
	  write_record ( fp, ptr );
      ptr = ptr->next;
    }
  return 1;
}

/* splice a record by a given letter */
FastaRecord* 
splice_c_record (const FastaRecord* src_record, const char letter)
{
  int i,num;
  FastaRecord * records;
  FastaRecord * cur_record;  
  char *src_seq;
  char **seqs;
  char *new_id;
  char *num_part;
  char *delimit = (char *)malloc(sizeof(char)*2);

  src_seq = (char *)strdup(src_record->fasta->seq);

  delimit[0] = letter, delimit[1] = '\0';
  seqs = splitstr( src_seq, &num, delimit);

  new_id = (char *)malloc(sizeof(char)*(strlen(src_record->fasta->id)+8));
  strcpy(new_id, src_record->fasta->id);
  strcat(new_id, "_part_1");
  records = new_record( new_id, seqs[0] );
  cur_record = records;
  for ( i=1 ; i<num ; i++ )
	{
	  /* new_id = old_id_part_i */
	  num_part = int2str(i+1);
	  new_id = (char *)malloc(sizeof(char)*(strlen(src_record->fasta->id)+8+strlen(num_part)));
	  strcpy(new_id, src_record->fasta->id);
	  strcat(new_id, "_part_");
	  strcat(new_id, num_part);
	  cur_record->next = new_record( new_id, seqs[i] );
	  cur_record = cur_record->next;
	  free(new_id);
	  free(num_part);
	}
  free (src_seq);
  return records;
}

/* splice a record by a given string */
FastaRecord* 
splice_s_record (const FastaRecord* src_record)
{
  return NULL;
}


/* Copy following two funcs from Eddy's squid-1.9 
 *
 * Function: StrReverse()
 * Date:     SRE, Thu Nov 20 10:54:52 1997 [St. Louis]
 * 
 * Purpose:  Returns a reversed version of s2, in s1.
 *           (s1 and s2 can be identical, to reverse in place)
 * 
 * Args:     s1 - allocated space for reversed string.
 *           s2 - string to reverse.
 *           
 * Return:   1.
 */                
int
StrReverse(char *s1, char *s2)
{
  int  len;
  int  pos;
  char c;
  
  len = strlen(s2);
  for (pos = 0; pos < len/2; pos++)
    {				/* swap ends */
      c             = s2[len-pos-1];
      s1[len-pos-1] = s2[pos];
      s1[pos]       = c;
    }
  if (len%2) { s1[pos] = s2[pos]; } /* copy middle residue in odd-len s2 */
  s1[len] = '\0';
  return 1;
}

/* Function: revcomp()
 *
 * Purpose:  Reverse complement fasta seq; store in comp.
 *           Can revcomp "in place" (revcomp(seq, seq)).
 *
 * Args:     comp  - destination for reverse complement of seq
 *           seq   - sequence to reverse complement
 *
 * Returns:  NULL on failure; or a (useless) pointer to comp.
 */
char *
revcomp(char *comp, char *seq)
{
  char *s;
  char  c;

  if (comp == NULL) return NULL;
  if (seq == NULL)  return NULL;

  StrReverse(comp, seq);
  for (s = comp; *s != '\0'; s++)
    {
      c = *s;
      c = toupper(c);
      switch (c) {
      case 'A': c = 'T'; break;
      case 'C': c = 'G'; break;
      case 'G': c = 'C'; break;
      case 'T': c = 'A'; break;
      case 'U': c = 'A'; break;
      case 'R': c = 'Y'; break;
      case 'Y': c = 'R'; break;
      case 'M': c = 'K'; break;
      case 'K': c = 'M'; break;
      case 'S': c = 'S'; break;
      case 'W': c = 'W'; break;
      case 'H': c = 'D'; break;
      case 'D': c = 'H'; break;
      case 'B': c = 'V'; break;
      case 'V': c = 'B'; break;
      default:  break;		/* anything else? leave it; it's prob a gap or an X */
      }
      if (islower((int) *s)) c = (char) tolower((int) c);
      *s = c;
    }
  return comp;
}


/* strip out \n */
static void 
remove_enter ( const char *seq ) 
{
  char *ptr_src, *ptr_trgt;
  ptr_src = strchr(seq, '\n');
  /* if \n is not found, no need to run this code. */
  if (ptr_src != NULL)
    {
      ptr_trgt = ptr_src;
      for (;*ptr_src; ptr_trgt++ ) 
		{
		  ptr_src++;
		  while (*ptr_src == '\n')
			ptr_src++;
		  *ptr_trgt = *ptr_src;
		}
    }
}

/* free mem */
void 
free_records (FastaRecord* target_record)
{
  FastaRecord *ptr,*tptr;
  ptr = target_record;
  while (ptr != NULL) 
    {
      tptr = ptr;
      ptr = ptr->next;
	  free(tptr->fasta->seq);
	  free(tptr->fasta->id);
      free(tptr->fasta);
      free(tptr);
    }
}
